import logging
import os

import streamlit as st
from llama_index.core import Settings
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)
'''
streamlit run class03/p3_3_3streamlit.py
'''


@st.cache_resource(show_spinner=False)
def load_data():
    # 1. 配置阿里云API（替换成您的API Key）
    print('初始化模型对象')
    llm_dash = DashScope(
        model="qwen-max",  # 使用通义千问Max模型 qwen-max qwen-turbo
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
        temperature=0.3,
        system_prompt="""作为一个虚拟对话助手，我将尝试与用户聊天。
              同时，在用户需要时，帮助用户提出问题并回答关于下一步的问题。
              仅根据提供的文档内容进行回答,如果话题与文档内容不相关,回答:抱歉这个话题我不擅长。
              """
    )

    # 2. 使用本地嵌入模型（节省成本）
    print('使用本地嵌入模型（节省成本）')
    embed_model_zh = HuggingFaceEmbedding(
        model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
        device="cpu"
    )
    with st.spinner(text="请稍等，正在处理。这可能需要1～2分钟。"):
        data_dir = 'D:/pythonRoom/llm_llama/data2'
        reader = SimpleDirectoryReader(input_dir=data_dir, recursive=True)
        docs = reader.load_data()
        # 配置全局设置
        Settings.embed_model = embed_model_zh
        Settings.llm = llm_dash
        index = VectorStoreIndex.from_documents(docs)
        return index


vector_index = load_data()

'''
streamlit run class03/p3_3_2streamlit.py
'''

st.session_state.messages = [
    {"role": "user", "content": "OPPO手机中最贵的是哪一款？"}
]

query_engine = vector_index.as_query_engine()
# 初始化一个ChatEngine实例并保存到session_state中
if "chat_engine" not in st.session_state.keys():
    chat_engine = CondenseQuestionChatEngine.from_defaults(
        query_engine=query_engine
    )
    st.session_state.chat_engine = chat_engine

# 初始化聊天消息历史列表
if "messages" not in st.session_state.keys():
    st.session_state.messages = [
        {"role": "assistant", "content": "欢迎。我来帮助你选购合适的手机。"}
    ]

# 获取用户输入并保存到聊天消息历史列表中
if prompt := st.chat_input("请输入你的问题"):
    st.write(f"用户：{prompt}")
    st.session_state.messages.append({"role": "user", "content": prompt})

# 如果是需要处理的用户消息，则调用LLM以获取响应结果
if st.session_state.messages[-1]["role"] != "assistant":
    with st.chat_message("assistant"):
        with st.spinner("处理中……"):
            response = st.session_state.chat_engine.chat(prompt)
            st.write(response.response)
            message = {"role": "assistant", "content": response.response}
            st.session_state.messages.append(message)
